import torch
import torch.nn as nn


class C1(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.conv = nn.Conv2d(cin, cout, kernel_size=3, padding=1)
        self.act = nn.SiLU()

    def forward(self, x):
        x = self.conv(x)
        return self.act(x)


class C2(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.con1 = C1(cin, cout)
        self.con2 = C1(cout, cout)
        self.con3 = C1(cin, cout)

    def forward(self, x):
        y = self.con1(x)
        y = self.con2(y)
        x = self.con3(x)
        return x + y


class C3(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.con1 = C1(cin, cout)
        self.con2 = C2(cout, cout)
        self.con3 = C1(cin, cout)

    def forward(self, x):
        y = self.con1(x)
        y = self.con2(y)
        x = self.con3(x)
        return torch.cat([x, y], dim=1)


class C4(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.con1 = C2(cin, cout)
        self.con2 = C2(cout, cout)
        self.con3 = C1(cin, cout)

    def forward(self, x):
        y = self.con1(x)
        y = self.con2(y)
        x = self.con3(x)
        return torch.cat([x, y], dim=1)

class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Sequential(
            nn.BatchNorm2d(3),
            C1(3, 16),
            C1(16, 32),
            nn.BatchNorm2d(32),
            C2(32, 16),
            C3(16, 8),
            C4(16, 4),
            nn.BatchNorm2d(8),
            C4(8, 4),
            C3(8, 4),
        )

        self.conv = nn.Sequential(
            nn.Conv2d(8, 4, kernel_size=5,padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(4, 1, kernel_size=5,padding=2),
            nn.MaxPool2d(2),
            nn.Sigmoid()
        )

        n=3750
        self.line=nn.Sequential(
            nn.Linear(n,640),nn.Sigmoid(),
            nn.Linear(640,320),nn.Sigmoid(),
            nn.Linear(320,320),nn.Sigmoid()
        )

    def forward(self, x):
        x = self.hidden(x)
        x = self.conv(x)
        x=x.view(-1,3750)
        x=self.line(x)
        return x.view(-1,64,5)


HEIGHT = 200
WIDTH = 300
CHANNEL = 3
if __name__ == "__main__":
    h, w, c = HEIGHT, WIDTH, CHANNEL
    if 0:
        model = torch.load("model.pt")
    else:
        model = MyNet()
    model.eval()
    inp = torch.rand(1, c, w, h)
    outp = model(inp)
    print(outp.shape)
    if True:
        torch.save(model, "model.pt")
        torch.onnx.export(model, inp, "./model.onnx")
